#include <sstream>
#include <iomanip>
#include <opencv2/dnn.hpp>

#include "base_util/utils.h"
#include "postprocess_op.h"
#include "preprocess_op.h"
#include "ocr_det.h"


namespace ai {


static void dump_tensor_attr(rknn_tensor_attr *attr) {
  printf("  index=%d, name=%s, n_dims=%d, dims=[%d, %d, %d, %d], n_elems=%d, size=%d, fmt=%s, type=%s, qnt_type=%s, "
          "zp=%d, scale=%f\n",
          attr->index, attr->name, attr->n_dims, attr->dims[0], attr->dims[1], attr->dims[2], attr->dims[3],
          attr->n_elems, attr->size, get_format_string(attr->fmt), get_type_string(attr->type),
          get_qnt_type_string(attr->qnt_type), attr->zp, attr->scale);
}
static inline float fast_exp(float x) {
  // return exp(x);
  union
  {
    uint32_t i;
    float f;
  } v;
  v.i = (12102203.1616540672 * x + 1064807160.56887296);
  return v.f;
}
static float DeQnt2F32(int8_t qnt, int zp, float scale) {
  return ((float)qnt - (float)zp) * scale;
}

RKOcrDetPredictor::RKOcrDetPredictor(RKPackPredictor* model, RKModelManager* manager, LogInfo *lg):
  RKPredictor(model, manager, lg) { }


int RKOcrDetPredictor::start_postprocess_thread_imp() {
  int ret = 0;
  std::pair<std::pair<int,int>,int> cur_block;
  // stop_flag2 = common::CameraGrabing;
  // while (stop_flag2 == common::CameraGrabing) {

  //   run_mtx2.lock();
  //   if (post_data2.empty()) {run_mtx2.unlock(); std::this_thread::sleep_for(std::chrono::milliseconds(3));continue;}
  //   cur_block = post_data2.front();
  //   post_data2.pop();
  //   run_mtx2.unlock();

  //   int loop_idx = cur_block.first.second;
  //   std::vector<BaseInfo*> det_infos;
  //   ret = det_hb_->PostProcess(det_infos, cur_block.first.first, this, cur_block.second);  //处理图片
  //   run_mtx3.lock();
  //   out_data3[loop_idx].push(det_infos);  //将处理后的数据返回到det_infos,加入队列中
  //   run_mtx3.unlock();
  // }
  // stop_flag2 = common::CameraOpened;
  return 0;
}

int RKOcrDetPredictor::RunDet(stream::ImageBlob* blob, std::vector<BaseInfo*>& det_infos) {
  int ret;

  stream::ImageBlob img_blob(stream::ImageBlobMode_BGR);
  if (!mdl_rk->cfg->transforms->run(blob->img, img_blob, mdl_rk->cfg)) { 
    printf("transforms->run fail \n");
    return model_image_channels_check_error;
  }
  
  std::cout << "img_blob.img w:" << img_blob.img.cols << " h:" << img_blob.img.rows << std::endl;
  // cv::imwrite("img_blob.img.jpg",img_blob.img);
  // cv::imwrite("blobimg.jpg",blob->img);

  int64_t start_time = 0, end_time = 0;
  start_time = duration_cast<milliseconds>(system_clock::now().time_since_epoch()).count();

  rknn_tensor_attr output_attrs[mdl_rk->output_count];
  memset(output_attrs, 0, sizeof(output_attrs));
  for (int i = 0; i < mdl_rk->output_count; i++) {
    output_attrs[i].index = i;
    ret = rknn_query(mdl_rk->ctx, RKNN_QUERY_OUTPUT_ATTR, &(output_attrs[i]), sizeof(rknn_tensor_attr));
    // dump_tensor_attr(&(output_attrs[i]));
  } 

  // mdl_rk->inputs[0].buf = (void *)(img_blob.img.data);
  mdl_rk->inputs[0].buf = (void *)(img_blob.img.data);
  rknn_inputs_set(mdl_rk->ctx, mdl_rk->input_count, &mdl_rk->inputs[0]);

  rknn_core_mask core_mask = RKNN_NPU_CORE_0;
  if(mdl_rk->cfg->gpu_id==1) { core_mask = RKNN_NPU_CORE_1; }
  if(mdl_rk->cfg->gpu_id==2) { core_mask = RKNN_NPU_CORE_2; }
  ret = rknn_set_core_mask(mdl_rk->ctx, core_mask);
  ret = rknn_run(mdl_rk->ctx, NULL);
  if (ret != 0) {printf("rknn_run fail %d", ret); return -1;}
  ret = rknn_outputs_get(mdl_rk->ctx, mdl_rk->output_count, &mdl_rk->outputs[0], NULL);
  if (ret != 0) {printf("rknn_outputs_get fail %d", ret); return -1;}
  end_time = duration_cast<milliseconds>(system_clock::now().time_since_epoch()).count();
  std::cout << "end_time - start_time:" << end_time - start_time << std::endl;

  // 后处理部分
  std::vector<float> out_scales;
  std::vector<int32_t> out_zps;
  for (int i = 0; i < mdl_rk->output_count; ++i) {
    out_scales.push_back(output_attrs[i].scale);
    out_zps.push_back(output_attrs[i].zp);
  }

  GetConvDetectionResult(img_blob, output_attrs, out_zps,  out_scales, det_infos);
  
  // for(int i=0;i<det_infos.size();i++)
  // {
  //   DetInfo* di = (DetInfo*)det_infos[i];
  //   cv::line(padimg,Point((int)di->poses[0].x,(int)di->poses[0].y),Point((int)di->poses[1].x,(int)di->poses[1].y),Scalar(255, 0, 255),2);
  //   cv::line(padimg,Point((int)di->poses[1].x,(int)di->poses[1].y),Point((int)di->poses[2].x,(int)di->poses[2].y),Scalar(255, 0, 255),2);
  //   cv::line(padimg,Point((int)di->poses[2].x,(int)di->poses[2].y),Point((int)di->poses[3].x,(int)di->poses[3].y),Scalar(255, 0, 255),2);
  //   cv::line(padimg,Point((int)di->poses[3].x,(int)di->poses[3].y),Point((int)di->poses[0].x,(int)di->poses[0].y),Scalar(255, 0, 255),2);
  // }
  // imwrite("draw.jpg",padimg);

  rknn_outputs_release(mdl_rk->ctx, mdl_rk->output_count, &mdl_rk->outputs[0]);

  return 0;

}

int RKOcrDetPredictor::GetConvDetectionResult(stream::ImageBlob& img_blob, rknn_tensor_attr* output_attrs, std::vector<int> &qnt_zp, std::vector<float> &qnt_scale, std::vector<BaseInfo*>& det_infos)
{

    int8_t *pblob = (int8_t *)mdl_rk->outputs[0].buf;
    float pf[output_attrs[0].n_elems]={0};
    for(int i=0;i<output_attrs[0].n_elems;i++)
    {
        pf[i]=DeQnt2F32(pblob[i], qnt_zp[0], qnt_scale[0]);
    }

    printf("[liufeng debug] h:%d w:%d\n",output_attrs[0].dims[2],output_attrs[0].dims[3]);
    uchar* bitmap=new uchar[output_attrs[0].dims[2]*output_attrs[0].dims[3]];

    for (int k = 0; k < output_attrs[0].dims[2] * output_attrs[0].dims[3]; k++)
    {
        if(pf[k]>0.3)
        {
            bitmap[k]=(int)(pf[k]*255);
        }
        else
        {
            bitmap[k]=0;
        }

    }

    cv::Mat imgmap=cv::Mat::zeros(output_attrs[0].dims[2],output_attrs[0].dims[3],CV_8UC1);
    for(int i=0;i<output_attrs[0].dims[2];i++)
    {
       uchar* data=imgmap.ptr(i);
       for(int j=0;j<output_attrs[0].dims[3];j++)
       {
          data[j]=bitmap[i*output_attrs[0].dims[3]+j];
       }
    }     
    // imwrite("imgmap.jpg",imgmap);

    delete [] bitmap;


    std::vector<std::vector<cv::Point>> poses;
    Plate_wordpos_det_baidu_model(imgmap,img_blob.ori_im_shape[1],img_blob.ori_im_shape[0],poses);


    std::vector<int> indices;
    for (int idx=0; idx < poses.size(); idx++) {indices.push_back(idx);}

    for (auto& idx : indices) {
      TextDetInfo* di = new TextDetInfo(poses[idx], 0.8);
      // di->poses=poses[idx];
      // printf("[liufeng debug] poses.size:%d \n",di->poses.size());
      // di->mask=imgmap;
      det_infos.emplace_back(di);
    }

    return 0;
}

int RKOcrDetPredictor::PostProcess(std::vector<BaseInfo*>& det_infos, 
                int cur_block,
                RKModelManager* mng, 
                int md_idx) {

  return 0;
}

void RKOcrDetPredictor::ParseTensor(float* data, int layer, 
  std::vector<cv::Rect>& boxs, std::vector<float>& scores, std::vector<int>& labels,
  std::vector<int> ori_shape, std::vector<int> feteature_shape, ModelConfig* cfg) {
    int num_classes = cfg->label_list.size();
    int stride = cfg->strides[layer];
    int num_pred = num_classes + 4 + 1;
    if (cfg->anchors_table.empty()) {
      spdlog::get("logger")->info("ERROR. RKOcrDetPredictor::ParseTensor anchors_table is empty.");
      return;
    }
    std::vector<std::vector<int>> &anchors = cfg->anchors_table[layer];
    float w_scale = ori_shape[1] * 1.0 / cfg->input_shape[1];
    float h_scale = ori_shape[0] * 1.0 / cfg->input_shape[0];
    // spdlog::get("logger")->info("num_classes: {}, stride: {}, num_pred: {}, w_scale: {}, h_scale: {}", num_classes, stride, num_pred, w_scale, h_scale);

    int anchor_num = anchors.size();
    for (int h = 0; h < feteature_shape[0]; h++) {
      for (int w = 0; w < feteature_shape[1]; w++) {
        for (int k = 0; k < anchor_num; k++) {
          float *cur_data = data + k * num_pred;
          // if (h == 0 && w == 0 ) {printf("%f, %f, %f, %f, %f, %f, %f, %f\n", cur_data[0],cur_data[1],cur_data[2],cur_data[3],cur_data[4],cur_data[5],cur_data[6],cur_data[7]);}
          
          int id = std::distance(cur_data + 5, std::max_element(cur_data + 5, cur_data + 5 + num_classes));
          double x1 = 1 / (1 + std::exp(-cur_data[4])) * 1;
          double x2 = 1 / (1 + std::exp(-cur_data[id + 5]));
          double confidence = x1 * x2;
          // if (1 ) {printf("id:%d, 4:%f, id+5:%f, x1:%f, x2:%f, confidence:%f\n", id, cur_data[4], cur_data[id + 5], x1, x2, confidence);}
          if (confidence < cfg->draw_threshold) {continue;}
          // spdlog::get("logger")->info("confidence:{}", confidence);

          double cx = ((1.0 / (1.0 + std::exp(-cur_data[0]))) * 2 - 0.5 + w) * stride;
          double cy = ((1.0 / (1.0 + std::exp(-cur_data[1]))) * 2 - 0.5 + h) * stride;
          double ww = std::pow((1.0 / (1.0 + std::exp(-cur_data[2]))) * 2, 2) * anchors[k][0];
          double hh = std::pow((1.0 / (1.0 + std::exp(-cur_data[3]))) * 2, 2) * anchors[k][1];

          double xmin = (cx - ww / 2.0)* w_scale;
          double ymin = (cy - hh / 2.0)* h_scale;
          double xmax = (cx + ww / 2.0)* w_scale;
          double ymax = (cy + hh / 2.0)* h_scale;

          if (xmax <= 0 || ymax <= 0) { continue; }
          if (xmin > xmax || ymin > ymax) { continue; }
          if (xmin < 0) { xmin = 0;}
          if (ymin < 0) { ymin = 0;}
          if (xmax >= ori_shape[1]) { xmax = ori_shape[1]-1;}
          if (ymax >= ori_shape[0]) { ymax = ori_shape[0]-1;}

          boxs.emplace_back(cv::Rect(xmin, ymin, xmax-xmin, ymax-ymin));
          scores.emplace_back((float)confidence);
          labels.emplace_back(id);
        }
        data = data + num_pred * anchors.size();
      }
    }
}


// int RKOcrDetPredictor::get_tensor_hw(hbDNNTensor* tensor, int *height, int *width) {
//   int h_index = 0;
//   int w_index = 0;
//   if (tensor->properties.tensorLayout == HB_DNN_LAYOUT_NHWC) {
//     h_index = 1;
//     w_index = 2;
//   }
//   else if (tensor->properties.tensorLayout == HB_DNN_LAYOUT_NCHW) {
//     h_index = 2;
//     w_index = 3;
//   }
//   else { return -1; }
//   *height = tensor->properties.validShape.dimensionSize[h_index];
//   *width = tensor->properties.validShape.dimensionSize[w_index];
//   return 0;
// }

void RKOcrDetPredictor::output_box(
    std::vector<float>& output_data,
    cv::Mat& img,
    stream::ImageBlob& img_info,
    std::vector<BaseInfo*>& det_infos,
    ModelConfig* cfg,
    const std::vector<float>& bbox) 
{
  // if (log_ifo->log_level_4) spdlog::get("logger")->info("1.2.0.4.0 rec_img");

  // int rh = 1;
  // int rw = 1;
  // std::vector<cv::Rect> boxs;
  // std::vector<float> scores;
  // std::vector<int> labels;
  // float nms_threshold;

  // if (cfg->arch == "YOLO_TORCH") {
  //   nms_threshold = 0.6;
  //   float w_ratio = img_info.ori_im_shape[1] * 1.0 / img_info.new_im_shape[1];
  //   float h_ratio = img_info.ori_im_shape[0] * 1.0 / img_info.new_im_shape[0];
  //   // log_tensor_shape("ori_im_shape", img_info->ori_im_shape);
  //   // log_tensor_shape("new_im_shape", img_info->new_im_shape);

  //   int class_num = cfg->label_list.size();
  //   int score_size = output_data.size() / class_num;
  //   int box_size = bbox.size() / 4;
  //   if (score_size != box_size) { return;} 

  //   for (int i = 0; i < box_size; i++) {
  //     auto it = std::max_element(output_data.begin() + i * class_num, output_data.begin() + (i+1) * class_num);
  //     if (*it < cfg->draw_threshold) { continue; }
  //     int w = bbox[i * 4 + 2] * img_info.ori_im_shape[1];
  //     int h = bbox[i * 4 + 3] * img_info.ori_im_shape[0];
  //     int x = bbox[i * 4 + 0] * img_info.ori_im_shape[1] - w / 2;
  //     int y = bbox[i * 4 + 1] * img_info.ori_im_shape[0] - h / 2;

  //     if (x < 0) { w += x; x = 0;}
  //     if (x + w > img.cols) { w = img.cols - x;}                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                           
  //     if (y < 0) { h += y; y = 0;}
  //     if (y + h > img.rows) { h = img.rows - y;}
  //     boxs.push_back(cv::Rect(x,y,w,h));
  //     scores.push_back(*it);
  //     labels.push_back(std::distance(output_data.begin() + i * class_num, it));
  //   }
  // }
  // else if (cfg->arch == "YOLOV5") {
  //   nms_threshold = 0.6;
  //   // float w_ratio = img_info->ori_im_shape[1] * 1.0 / img_info->new_im_shape[1];
  //   // float h_ratio = img_info->ori_im_shape[0] * 1.0 / img_info->new_im_shape[0];

  //   int class_num = cfg->label_list.size();
  //   int step = class_num+5;
  //   int box_size = output_data.size() / step;
  //   if (log_ifo->log_level_4) spdlog::get("logger")->info("1.2.0.4.0 class_num:{}, step:{} box_size:{}", class_num, step, box_size);

  //   for (int i = 0; i < box_size; i++) {
  //     auto it = std::max_element(output_data.begin() + i*step+5, output_data.begin() + i*step+step);
  //     float score = *it * output_data[i*step + 4];

  //     if (score < cfg->draw_threshold) { continue; }
  //     // if (score > 0.8) {
  //     //   std::cout << output_data[i*step + 0] * w_ratio<< "  " << output_data[i*step + 1]*h_ratio << "  " << output_data[i*step + 2] * w_ratio << "  " << output_data[i*step + 3] *h_ratio<< "  " << output_data[i*step + 4] << "  " << output_data[i*step + 5] << "  " << output_data[i*step + 6] << "  " << output_data[i*step + 7] << "  " << output_data[i*step + 8] << "  " << output_data[i*step + 9] << "  " << output_data[i*step + 10] << "  " << output_data[i*step + 11]<< std::endl;
  //     // }
  //     // std::cout << "score: " << score << std::endl;

  //     float w = output_data[i*step + 2] / img_info.scale;
  //     float h = output_data[i*step + 3] / img_info.scale;
  //     float x = output_data[i*step + 0] / img_info.scale - w/2;
  //     float y = output_data[i*step + 1] / img_info.scale - h/2;
  //     // std::cout << "(" << x << "," << y << "," << w<<","<<h<<") "  << std::endl;

  //     if (x < 0) { w += x; x = 0;}
  //     if (x + w > img.cols) { w = img.cols - x;}                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                           
  //     if (y < 0) { h += y; y = 0;}
  //     if (y + h > img.rows) { h = img.rows - y;}
  //     boxs.push_back(cv::Rect(x,y,w,h));
  //     scores.push_back(score);
  //     labels.push_back(std::distance(output_data.begin() + i*step+5, it));
  //   }
  // }
  // else {
  //   nms_threshold = 0.3;
  //   int box_size = output_data.size() / 6;
  //   for (int j = 0; j < box_size; ++j) {
  //     // Class id
  //     int class_id = static_cast<int>(round(output_data[0 + j * 6]));
  //     // Confidence score
  //     float score = output_data[1 + j * 6];
  //     int xmin = (output_data[2 + j * 6] * rw);
  //     int ymin = (output_data[3 + j * 6] * rh);
  //     int xmax = (output_data[4 + j * 6] * rw);
  //     int ymax = (output_data[5 + j * 6] * rh);

  //     if (score > cfg->draw_threshold && class_id > -1) {
  //       if (xmin < 0) { xmin = 0;}
  //       if (ymin < 0) { ymin = 0;}
  //       if (xmax > img.cols) { xmax = img.cols;}
  //       if (ymax > img.rows) { ymax = img.rows;}
  //       int wd = xmax - xmin;
  //       int hd = ymax - ymin;

  //       boxs.push_back(cv::Rect(xmin, ymin, wd, hd));
  //       scores.push_back(score);
  //       labels.push_back(class_id);
  //     }
  //   }
  // }
  // if (log_ifo->log_level_4) spdlog::get("logger")->info("1.2.0.4.1 boxs:{}", boxs.size());

  // std::vector<int> indices;
  // cv::dnn::NMSBoxes(boxs, scores, cfg->draw_threshold, nms_threshold, indices, 100);
  // for (auto& idx : indices) {
  //   DetInfo* di = new DetInfo(scores[idx], labels[idx], boxs[idx], cfg->label_list[labels[idx]]);
  //   det_infos.push_back(di);
  // }
  // if (log_ifo->log_level_4) spdlog::get("logger")->info("1.2.0.4.2 det_infos:{}", det_infos.size());

}

void RKOcrDetPredictor::output_box(
    std::vector<float>& output_data,
    std::vector<cv::Mat>& imgs,
    std::vector<stream::ImageBlob>& img_info,
    std::vector<std::vector<DetInfo>>& det_infos,
    ModelConfig* cfg) 
{
  int rh = 1;
  int rw = 1;

  // int total_size = output_data.size() / 6;
  // if (total_size <= 0) {
  //   det_infos.resize(imgs.size());
  //   return ;
  // }
  // int batch_idx = -1;
  // float last_score = 0.;
  // std::vector<DetInfo> det_info;
  // for (int j = 0; j < total_size; ++j) {
  //   // Class id
  //   int class_id = static_cast<int>(round(output_data[0 + j * 6]));
  //   // Confidence score
  //   float score = output_data[1 + j * 6];
  //   if (score > last_score) {
  //     if (j != 0) { det_infos.push_back(det_info); }
  //     det_info.clear();
  //     batch_idx++;
  //     if (cfg->arch == "SSD" || cfg->arch == "Face") {
  //       rh = imgs[batch_idx].rows;
  //       rw = imgs[batch_idx].cols;
  //     }
  //   }
  //   last_score = score;

  //   int xmin = (output_data[2 + j * 6] * rw);
  //   int ymin = (output_data[3 + j * 6] * rh);
  //   int xmax = (output_data[4 + j * 6] * rw);
  //   int ymax = (output_data[5 + j * 6] * rh);
  //   if (log_ifo->log_level_4) {
  //   }

  //   if (score > cfg->draw_threshold && class_id > -1) {
  //     if (xmin < 0) { xmin = 0;}
  //     if (ymin < 0) { ymin = 0;}
  //     if (xmax > imgs[batch_idx].cols) { xmax = imgs[batch_idx].cols;}
  //     if (ymax > imgs[batch_idx].rows) { ymax = imgs[batch_idx].rows;}
  //     int wd = xmax - xmin;
  //     int hd = ymax - ymin;

  //     cv::Rect box = cv::Rect(xmin, ymin, wd, hd);

  //     det_info.push_back(DetInfo(class_id,score, box));
  //   }
  // }
  // det_infos.push_back(det_info); 
}

}  // namespace hb
